# 
# - Step 3 (freq.py): collecting sucess-frequencies for the generated Insights
# 	- For each question + reasoning history, generate a reasoning-situation
# 	- Use the embedding of the reasoning-situation to find K most similar Insights
# 	- Randomly select one Insight as example to generate an insight for the curren the reasoning-situation
# 	- Use Insights generated this way to generate the complete reasoning for the question
# 	- If the reasoning results is a correct answer, increase the frequency of the referred Insights, and save the embedding of the embedding of the reasoning-situation and that of the referred Insight into a similarity training set
# 	- Repeat the above steps, and finally use the frequencies to trim the library and keep only the most freqently successful Insights

from exp import data, run, math_eval
from com import tbys, emb
from datetime import datetime
import math, os

freq_path = '../OUT/freq'
examples_path = '../OUT/created-examples'

def collect_freq(test_set, num_exp=12):
	reasoner_1 = tbys.Reasoner(*tbys.load_data_1(test_set))
	dataset = data.get_data(test_set)
	for i in range(num_exp):
		outfile = f'{freq_path}/step3_{test_set}_freq_{i+1:03d}.jsonl'
		print(f'# {i+1} {outfile} {datetime.now()}')
		ex = run.run_exp(dataset, reasoner_1.reason, num_samples=10, max_retries=30, outfile=outfile)
		if ex: break

def eval_freq(test_set, *paths, dump_freq=False):
	example_count = {}
	num_correct_ = 0
	num_total_ = 0
	for path in paths:
		print(path)
		num_correct = 0
		res = emb.load_jsonl(path)
		for r in res:
			gt = r['solution']
			ans = r['pred'][1]
			eq, gt_, ans_ = math_eval.check(gt, ans)
			if eq:
				num_correct += 1
			if 'steps' in r:
				for step in r['steps']:
					if step['examples'] is None: continue
					for example_id, step_num in step['examples']:
						if example_id not in example_count:
							example_count[example_id] = [0,0]
						if eq:
							example_count[example_id][1] += 1
						else:
							example_count[example_id][0] += 1
		print(f'Acc: {num_correct}/{len(res)} = {100*num_correct/len(res):.2f}%')
		num_correct_ += num_correct
		num_total_ += len(res)
	if len(paths) > 1:
		print(f'Avg Acc: {num_correct_}/{num_total_} = {100*num_correct_/num_total_:.2f}%')
	if dump_freq:
		emb.dump_jsonl([example_count], f'{freq_path}/step3_{test_set}_freq.jsonl')
	return num_correct_/num_total_

def reorder_examples(test_set, unique_example=True):
	count = emb.load_jsonl(f'{freq_path}/step3_{test_set}_freq.jsonl')[0]
	count = [[int(n),[w,r]] for n,(w,r) in count.items() if r+w>9]
	total_ = 0
	for n,[w,r] in count:
		total_ += w + r
	def print_avg(count_):
		avg_w = []
		avg_f = []
		min_w = 1
		total = 0
		for n,[w,r] in count_:
			total += w + r
			min_w = min(min_w, r/(r+w))
			avg_w.append(r/(r+w))
			avg_f.append(r+w)
		print('min_win', min_w)
		print('avg_win', sum(avg_w)/len(avg_w))
		print('avg_freq', sum(avg_f)/len(avg_f))
		print('coverage', total/total_)
	def sort_key(x):
		n,[w,r] = x
		# return r/(r+w)
		return r/(r+w) * math.log(r+w)
	count.sort(key=sort_key, reverse=True)
	# load files
	example_paths = [f'{examples_path}/step1_{test_set}.jsonl']
	emb_path = f'{examples_path}/step2_{test_set}_emb.jsonl'
	examples = emb.load_examples(example_paths)
	embs = emb.load_jsonl(emb_path)
	# make example unique
	unique_pids = set()
	if unique_example:
		count, count_0 = [], count
		for n,[w,r] in count_0:
			pid = examples[n]['pid']
			if pid not in unique_pids:
				unique_pids.add(pid)
				count.append([n,[w,r]])
	sizes = [5000, 2500, 1000, 500, 250, 100, 50]
	for size in sizes:
		if size <= len(count):
			print('>> size', size)
			print_avg(count[:size])
	# produce re-order files	
	examples2 = [examples[n] for n,(w,r) in count]
	embs2 = [embs[n] for n,(w,r) in count]
	emb.dump_jsonl(examples2, f'{freq_path}/examples_{test_set}.jsonl')
	emb.dump_jsonl(embs2, f'{freq_path}/emb_{test_set}.jsonl')

def step3(test_set='MATH500'):
	if not os.path.isdir(freq_path):
		os.mkdir(freq_path)
	collect_freq(test_set)
	paths = [os.path.join(freq_path, f) for f in os.listdir(freq_path) if f.startswith(f'step3_{test_set}_freq_') and f.endswith('.jsonl')]
	eval_freq(test_set, *paths, dump_freq=True)
	reorder_examples(test_set)

if __name__ == '__main__':
	step3()
